package com.ada.spark.sparksql

import org.apache.spark.SparkConf
import org.apache.spark.rdd.RDD
import org.apache.spark.sql.expressions.{MutableAggregationBuffer, UserDefinedAggregateFunction}
import org.apache.spark.sql.types._
import org.apache.spark.sql.{Dataset, Row, SparkSession}

/**
  * 强类型的Dataset和弱类型的DataFrame都提供了相关的聚合函数， 如 count()，countDistinct()，avg()，max()，min()。
  * 除此之外，用户可以设定自己的自定义聚合函数。
  */
object SparkSql06_UDAF {

    def main(args: Array[String]) {

        //创建配置对象
        val conf: SparkConf = new SparkConf().setAppName("SparkSql06_UDAF").setMaster("local[*]")

        //创建SparkSql的环境对象
        val spark: SparkSession = SparkSession.builder().config(conf).getOrCreate();

        //进行转换之前，需要引入隐式转换规则
        //这里的spark不是包名的含义，是SparkSession对象的名字
        import spark.implicits._

        //创建RDD
        val rdd: RDD[(Int, String, Int)] = spark.sparkContext.makeRDD(List((1, "Jack", 20), (2, "Bob", 22), (3, "Lucy", 28)))

        //RDD -> DataSet
        val userRDD: RDD[User] = rdd.map {
            case (id, name, age) => {
                User(id, name, age)
            }
        }

        val userDS: Dataset[User] = userRDD.toDS()

        //创建临时视图
        userDS.createOrReplaceTempView("user")

        //注册自定义聚合函数
        spark.udf.register("myAverage", MyAverage)

        //使用自定义聚合函数
        spark.sql("select myAverage(age) from user").show()

        //释放资源
        spark.stop()

    }

}

/**
  * 声明自定义聚合函数(弱类型)
  */
object MyAverage extends UserDefinedAggregateFunction {

    // 聚合函数输入参数的数据类型
    def inputSchema: StructType = StructType(StructField("age", LongType) :: Nil)

    // 聚合缓冲区中值的数据类型
    def bufferSchema: StructType = {
        StructType(StructField("sum", LongType) :: StructField("count", LongType) :: Nil)
    }

    // 返回值的数据类型
    def dataType: DataType = DoubleType

    // 对于相同的输入是否一直返回相同的输出。
    def deterministic: Boolean = true

    // 初始化
    def initialize(buffer: MutableAggregationBuffer): Unit = {
        // sum
        buffer(0) = 0L
        // count
        buffer(1) = 0L
    }

    // 相同Execute间的数据合并。
    def update(buffer: MutableAggregationBuffer, input: Row): Unit = {
        if (!input.isNullAt(0)) {
            // sum
            buffer(0) = buffer.getLong(0) + input.getLong(0)
            // count
            buffer(1) = buffer.getLong(1) + 1
        }
    }

    // 不同Execute间的数据合并
    def merge(buffer1: MutableAggregationBuffer, buffer2: Row): Unit = {
        // sum
        buffer1(0) = buffer1.getLong(0) + buffer2.getLong(0)
        // count
        buffer1(1) = buffer1.getLong(1) + buffer2.getLong(1)
    }

    // 计算最终结果
    def evaluate(buffer: Row): Double = {
        //sum/count
        buffer.getLong(0).toDouble / buffer.getLong(1)
    }

}


